# person_track_udp.py
import cv2
import socket
import time

# 模型配置
prototxt = "models/deploy.prototxt"
model = "models/mobilenet_iter_73000.caffemodel"

CLASSES = ["background", "aeroplane", "bicycle", "bird", "boat",
           "bottle", "bus", "car", "cat", "chair", "cow", "diningtable",
           "dog", "horse", "motorbike", "person", "pottedplant",
           "sheep", "sofa", "train", "tvmonitor"]

net = cv2.dnn.readNetFromCaffe(prototxt, model)

# 摄像头
cap = cv2.VideoCapture(1, cv2.CAP_V4L2)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

if not cap.isOpened():
    print("[ERROR] 无法打开摄像头")
    exit()

# UDP
UDP_IP = "127.0.0.1"
UDP_PORT = 5005
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

# 相机参数
H_FOV_DEG = 60.0
V_FOV_DEG = 45.0
IMG_W, IMG_H = 640, 480
deg_per_px_h = H_FOV_DEG / IMG_W
deg_per_px_v = V_FOV_DEG / IMG_H
print(f"[INFO] deg_per_px_h={deg_per_px_h:.6f}, deg_per_px_v={deg_per_px_v:.6f}")

# 参数
CONF_THRESH = 0.5
MIN_MOVE_DEG = 0.5
SEND_INTERVAL = 0.05
last_send_time = 0
k_speed = 5.0
MAX_SPEED = 120.0

print("[INFO] 开始检测与跟踪。按 q 退出。")

while True:
    ret, frame = cap.read()
    if not ret:
        print("[ERROR] 读取帧失败")
        break

    (h, w) = frame.shape[:2]
    blob = cv2.dnn.blobFromImage(cv2.resize(frame, (300, 300)),
                                 0.007843, (300, 300), 127.5)
    net.setInput(blob)
    detections = net.forward()

    best_conf = 0.0
    best_box = None

    for i in range(detections.shape[2]):
        confidence = float(detections[0, 0, i, 2])
        if confidence > CONF_THRESH:
            idx = int(detections[0, 0, i, 1])
            if CLASSES[idx] != "person":
                continue
            box = detections[0, 0, i, 3:7] * [w, h, w, h]
            (startX, startY, endX, endY) = box.astype("int")
            if confidence > best_conf:
                best_conf = confidence
                best_box = (startX, startY, endX, endY)

    if best_box is not None:
        (sx, sy, ex, ey) = best_box
        cx = (sx + ex) / 2.0
        cy = (sy + ey) / 2.0

        cv2.rectangle(frame, (sx, sy), (ex, ey), (0, 255, 0), 2)
        label = f"person: {best_conf:.2f}"
        cv2.putText(frame, label, (sx, sy - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        cv2.circle(frame, (int(cx), int(cy)), 4, (0, 0, 255), -1)

        img_cx, img_cy = w / 2.0, h / 2.0
        dx_px, dy_px = cx - img_cx, cy - img_cy

        # === X轴取反，Y轴保持 ===
        pan_deg = -dx_px * deg_per_px_h
        tilt_deg = dy_px * deg_per_px_v

        now = time.time()
        if now - last_send_time >= SEND_INTERVAL:
            send_pan, send_tilt = pan_deg, tilt_deg

            if abs(send_pan) < MIN_MOVE_DEG:
                send_pan = 0.0
            if abs(send_tilt) < MIN_MOVE_DEG:
                send_tilt = 0.0

            if send_pan != 0.0 or send_tilt != 0.0:
                speed = min(MAX_SPEED, max(10.0, k_speed * max(abs(send_pan), abs(send_tilt))))
                cmd = f"MOVE {send_pan:.3f} {send_tilt:.3f} {speed:.1f}"
                try:
                    sock.sendto(cmd.encode('utf-8'), (UDP_IP, UDP_PORT))
                    print(f"[SEND] {cmd}")
                except Exception as e:
                    print(f"[ERROR] 发送 UDP 失败: {e}")
                last_send_time = now

    cv2.line(frame, (int(w/2)-10, int(h/2)), (int(w/2)+10, int(h/2)), (255,0,0), 1)
    cv2.line(frame, (int(w/2), int(h/2)-10), (int(w/2), int(h/2)+10), (255,0,0), 1)

    cv2.imshow("Person Tracking", frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()
sock.close()
print("[INFO] 退出程序。")
